import os
import sys
import warnings
from functools import partial
import pathlib
import skrf

from skrf.network import two_port_reflect, four_oneports_2_twoport
from skrf.calibration import OnePort, TwoPortOnePath, EnhancedResponse
from skrf.media import DefinedGammaZ0


# global variables
cal_dir = pathlib.Path(sys.argv[1])
measurement_dir = pathlib.Path(sys.argv[2])
export_dir = pathlib.Path(sys.argv[3])


def get_standards(medium, port=1):
    if port != 1 or port != 2:
        raise ValueError

    # modeled SOLT
    short_model = skrf.Network(cal_dir / "model" / "short.s1p")
    if not short_model.exists():
        short_model = medium.short()

    open_model  = skrf.Network(cal_dir / "model" / "open.s1p")
    if not open_model.exists():
        open_model = medium.open()

    load_model  = skrf.Network(cal_dir / "model" / "load.s1p")
    if not load_model.exists():
        load_model = medium.match()

    thru_model = skrf.Network(cal_dir / "model" / "thru.s2p")
    if not thru_model.exists():
        thru_model = medium.thru(nports=2)

    # measured SOLT
    if (cal_dir / "actual").exists():
        cal_dir = cal_dir / "actual"

    short_actual = skrf.Network(cal_dir / "short.s1p")
    open_actual  = skrf.Network(cal_dir / "open.s1p")
    load_actual  = skrf.Network(cal_dir / "load.s1p")
    thru_actual  = skrf.Network(cal_dir / "thru.s2p")

    if port == 1:
        modeled = [short_model,  open_model,  load_model]
        actual  = [short_actual, open_actual, load_actual]
    elif port == 2:
        modeled = [short_model,  open_model,  load_model]
        actual  = [short_actual, open_actual, load_actual]

        for idx, network in enumerate(modeled):
            modeled[idx] = skrf.network.two_port_reflect(network, network)
        for idx, network in enumerate(actual):
            actual[idx] = skrf.network.two_port_reflect(network, network)

        # Thru is a two-port network
        modeled.append(thru_model)
        actual.append(thru_actual)

    freq = open_actual.freq
    return frequency, modeled, actual


def replace_filename_segment(f, orig, replace):
    suffix = f.suffix
    name_segments = f.stem.split("-")
    if name_segments[-1] != orig:
        raise ValueError

    if replace:
        name_segments[-1] = replace
    else:
        del name_segments[-1]

    return f.with_name("-".join(name_segments) + suffix)


def one_half_ports_to_two_ports(network):
    s11 = network.s11
    s22 = s11
    try:
        s21 = network.s21
        s12 = s21
    except:
        s21 = skrf.Network(frequency=network.f, s=[0] * len(network.f))
        s12 = s21
    return four_oneports_2_twoport(s11, s12, s21, s22)


def two_port_one_path(dut_forward, dut_backward=None, export_file='/dev/null'):
    line = DefinedGammaZ0(frequency=frequency)
    frequency, ideals, measured = get_standards(line)

    frequency = skrf.frequency.overlap_freq(frequency, dut_forward.frequency)
    dut_forward = dut_forward.interpolate(frequency)
    dut_backward = dut_backward.interpolate(frequency)



    #global short_std, open_std, load_std, thru_std

    # pull overlapping frequency information from measurements,
    # and interpolate the overlapping parts.
    frequency = skrf.frequency.overlap_freq(open_std.frequency, dut_forward.frequency)
    dut_forward = dut_forward.interpolate(frequency)
    dut_backward = dut_backward.interpolate(frequency)
    _open_std = open_std.interpolate(frequency)
    _short_std = short_std.interpolate(frequency)
    _load_std = load_std.interpolate(frequency)
    _thru_std = thru_std.interpolate(frequency)

    # convert 1.5-port measurements to 2-port data by faking data
    short_std_2 = one_half_ports_to_two_ports(_short_std)
    open_std_2 = one_half_ports_to_two_ports(_open_std)
    load_std_2 = one_half_ports_to_two_ports(_load_std)
    thru_std_2 = one_half_ports_to_two_ports(_thru_std)


    # apply calibration to DUT
    cal = TwoPortOnePath(measured=measured, ideals=ideals, n_thrus=1)

    if not dut_backward:
        # Enhanced Response
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            dut_correction = cal.apply_cal(dut_forward)
    else:
        # Full correction or FakeFlip
        dut_correction = cal.apply_cal((dut_forward, dut_backward))
    dut_correction.write_touchstone(export_file)
    print(export_file)


def one_port(dut, export_file='/dev/null'):
    # pull overlapping frequency information from measurements,
    # and interpolate the overlapping parts.
    frequency = skrf.frequency.overlap_freq(open_std.frequency, dut.frequency)

    dut = dut.interpolate(frequency)
    _open_std = open_std.interpolate(frequency)
    _short_std = short_std.interpolate(frequency)
    _load_std = load_std.interpolate(frequency)

    # List of 'ideal' responses of the calibration standards
    # and corresponding measurements to the 'ideals'.
    line = DefinedGammaZ0(frequency=frequency)
    ideals, measured = get_standards(line)

    # the Calibration object
    cal = OnePort(measured=measured, ideals=ideals)

    # apply calibration to DUT
    dut_correction = cal.apply_cal(dut)
    dut_correction.write_touchstone(export_file)
    print(export_file)
    

execution_list = []

# Two-Port Calibration
for f in measurement_dir.glob("*.s2p"):
    if f.name.endswith("-forward.s2p"):
        # Full Calibration
        f_backward = replace_filename_segment(f, "forward", "backward")
        output = export_dir / replace_filename_segment(f, "forward", "").name

        if f_backward.exists():
            s2p_forward = skrf.Network(f)
            s2p_backward = skrf.Network(f_backward)
            execution_list += [partial(two_port_one_path, s2p_forward, s2p_backward, export_file=output)]
        else:
            raise FileNotFoundError("%s calibration failed: %s not found!" % (f.name, f_backward.name))
    elif f.name.endswith("-backward.s2p"):
        pass
    elif f.name.endswith("-fakeflip.s2p"):
        # FakeFlip Calibration
        s2p_forward = skrf.Network(f)
        output = export_dir / replace_filename_segment(f, "fakeflip", "").name
        execution_list += [partial(two_port_one_path, s2p_forward, s2p_forward, export_file=output)]
    else:
        # Enhanced Response
        s2p_forward = skrf.Network(f)
        output = export_dir / f.name
        execution_list += [partial(two_port_one_path, s2p_forward, export_file=output)]

# One-Port Calibration
for f in measurement_dir.glob("*.s1p"):
    s1p = skrf.Network(f)
    output = export_dir / f.name
    execution_list += [partial(one_port, s1p, output)]

for task in execution_list:
    task()
